# -*- coding: UTF-8 -*-
"""HaaS506 DTU MQTT Module - 支持二级 topic 订阅和 /write 控制指令处理版
"""

import mqtt
import ujson
import utime as time
from driver import TIMER
import config
import mod_network
import mod_dispatch

# 全局变量
_client = None
_config = None
_connected = False
_running = False
_timer = None
_upload_results = {}
module_status = "STOPPED"

_stats = {
    "messages_sent": 0,
    "connection_attempts": 0,
    "last_connect_time": 0,
    "upload_failures": 0,
    "upload_success": 0,
    "result_received": 0,
    "last_reconnect_attempt": 0,
    "network_down_count": 0,
    "last_network_check": 0,
}


def safe_operation(fn, *args, **kwargs):
    try:
        return fn(*args, **kwargs)
    except Exception as e:
        print("[MQTT] Operation failed: {}".format(e))
        return None


def _on_message_callback(topic, payload):
    try:
        print("[MQTT] Received message on topic: {}".format(topic))
        message = ujson.loads(payload)
        print("[MQTT] call back from {},  Message: {}".format(topic, message))

        if topic == _config["topics"].get("result"):
            no = message.get("no", "unknown")
            success = message.get("success", False)
            timestamp = message.get("timestamp", 0)

            _stats["result_received"] += 1
            if success:
                _stats["upload_success"] += 1
                print("[MQTT] Upload SUCCESS for {}".format(no))
            else:
                print("[MQTT] Upload FAILED for {}".format(no))

            return timestamp

        elif topic == _config["topics"].get("control"):
            try:
                mod_dispatch.handle_action(message)
                print("[MQTT] Control action executed successfully")
            except Exception as control_error:
                print("[MQTT] Control action failed: {}".format(control_error))
                return 0
        return 0
    except Exception as e:
        print("[MQTT] Message callback error: {}".format(e))
        return 0


def _load_config():
    global _config
    _config = config.get_mqtt()
    if not _config:
        print("[MQTT] Failed to load config")
        return False
    print("[MQTT] Broker: {}:{}".format(_config["broker"], _config["port"]))
    return True


def _subscribe_all_topics():
    try:
        topics = _config.get("topics", {})
        for key, topic in topics.items():
            if not topic:
                continue
            # 使用on_subcribe回调后，订阅时不需要传递回调函数
            res = safe_operation(_client.subscribe, topic, _config["qos"])
            if isinstance(res, tuple):
                res_code = res[0]
            else:
                res_code = res
            if res_code == 0:
                print("[MQTT] Subscribed to topic '{}': {}".format(key, topic))
            else:
                print(
                    "[MQTT] Failed to subscribe '{}', code: {}".format(topic, res_code)
                )
        return True
    except Exception as e:
        print("[MQTT] Subscribe exception: {}".format(e))
        return False


def _connect():
    global _client, _connected, _stats

    _stats["connection_attempts"] += 1
    print("[MQTT] Connecting attempt #{}".format(_stats["connection_attempts"]))

    connect_timeout = 10
    start_time = time.time()
    while not mod_network.getNetFlag():
        if time.time() - start_time > connect_timeout:
            print("[MQTT] Network connection timeout, aborting connect")
            return False
        print("[MQTT] Waiting for network...")
        time.sleep(1)

    try:
        # 清理旧的客户端连接
        if _client:
            try:
                _client.disconnect()
            except:
                pass
            _client = None

        _client = mqtt.client(_config["client_id"])

        # 设置订阅回调函数，替代频繁的loop调用
        _client.on_subcribe(_on_message_callback)

        if _config.get("username") and _config.get("password"):
            _client.username_pw_set(_config["username"], _config["password"])

        port = int(_config.get("port", 1883))
        result = safe_operation(
            _client.connect, _config["broker"], port, _config["keepalive"]
        )
        if result != 0:
            print(
                "[MQTT] Connect failed, code: {}，_config: {}，_config".format(
                    result, _config
                )
            )
            _client = None
            return False

        _connected = True
        _stats["last_connect_time"] = time.time()
        print("[MQTT] Connected successfully")
        return _subscribe_all_topics()

    except Exception as e:
        print("[MQTT] Connect exception: {}".format(e))
        _client = None
        _connected = False
        return False


def is_connected():
    return _connected and _client is not None


def publish_data(data, topic=None):
    global _stats
    if not is_connected():
        print("[MQTT] Not connected")
        return False

    topic = topic or _config.get("topics", {}).get("data")
    if not topic:
        print("[MQTT] No publish topic configured")
        return False

    payload = ujson.dumps(data)
    result = safe_operation(_client.publish, topic, payload, _config["qos"])

    # 使用on_subcribe回调后，服务器响应会自动通过回调处理
    # 不需要手动调用loop

    if result is not None and result >= 0:
        _stats["messages_sent"] += 1
        print("[MQTT] Data published")
        return True
    else:
        _connected = False
        _stats["upload_failures"] += 1
        print(
            "[MQTT] Publish 1 failed {} times,result:{}".format(
                _stats["upload_failures"], result
            )
        )
        return False


def publish_status(data, topic=None):
    if not is_connected():
        print("[MQTT] Not connected")
        return False

    topic = topic or _config.get("topics", {}).get("status")
    if not topic:
        print("[MQTT] No publish topic configured")
        return False

    payload = ujson.dumps(data)
    result = safe_operation(_client.publish, topic, payload, _config["qos"])

    # 使用on_subcribe回调后，服务器响应会自动通过回调处理
    # 不需要手动调用loop

    if result is not None and result >= 0:
        _stats["messages_sent"] += 1
        print("[MQTT] Data published")
        return True
    else:
        _connected = False
        _stats["upload_failures"] += 1
        print("[MQTT] Publish 2 failed,result:{}".format(result))
        return False


def _timer_callback(args):
    global _connected

    try:
        current_time = time.time()

        if not mod_network.getNetFlag():
            # 网络异常计数和处理
            if current_time - _stats["last_network_check"] >= 1:  # 每秒检查一次
                _stats["network_down_count"] += 1
                _stats["last_network_check"] = current_time
                print("[MQTT] Network down ({})".format(_stats["network_down_count"]))

                # 网络异常超过30次（约30秒），尝试重新连接网络
                if _stats["network_down_count"] >= 30:
                    print(
                        "[MQTT] Network down too long, attempting network reconnect..."
                    )
                    _stats["network_down_count"] = 0  # 重置计数
                    try:
                        mod_network.connect_network()
                    except Exception as net_error:
                        print("[MQTT] Network reconnect failed: {}".format(net_error))

            _connected = False
            return
        else:
            # 网络恢复，重置计数
            if _stats["network_down_count"] > 0:
                print(
                    "[MQTT] Network recovered after {} checks".format(
                        _stats["network_down_count"]
                    )
                )
                _stats["network_down_count"] = 0

        if not is_connected():
            # 控制重连频率，避免频繁重连
            if current_time - _stats["last_reconnect_attempt"] >= 5:  # 至少间隔5秒
                print("[MQTT] Disconnected, retrying connect...")
                _stats["last_reconnect_attempt"] = current_time
                _connect()
        # 使用on_subcribe回调后，不需要频繁调用loop处理消息
        # 消息会通过回调函数自动处理
    except Exception as e:
        print("[MQTT] Timer callback error: {}".format(e))


def init():
    global module_status

    if not _load_config():
        module_status = "ERROR"
        return False

    module_status = "INITIALIZED"
    start()
    return True


def start():
    global _running, module_status, _timer

    if _running:
        return True

    _running = True
    module_status = "RUNNING"

    try:
        _timer = TIMER(0)
        # 提高定时器频率到500ms，增强消息处理实时性
        _timer.open(period=500, mode=TIMER.PERIODIC, callback=_timer_callback)
        print("[MQTT] Timer started with 500ms interval")
    except Exception as e:
        print("[MQTT] Timer error: {}".format(e))

    _connect()
    return True


def stop():
    global _running, _connected, _client, module_status, _timer

    _running = False
    _connected = False
    module_status = "STOPPED"

    if _timer:
        try:
            _timer.close()
        except:
            pass
        _timer = None

    if _client:
        safe_operation(_client.disconnect)
        _client = None

    print("[MQTT] Stopped")


def process_messages():
    """处理MQTT消息 - 使用on_subcribe回调后此函数已不需要"""
    # 消息现在通过on_subcribe回调自动处理
    return is_connected()


def check_control_messages():
    """检查控制指令消息 - 使用on_subcribe回调后此函数已不需要"""
    # 控制消息现在通过on_subcribe回调自动处理
    return is_connected()
